{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "56e7abc3",
   "metadata": {},
   "source": [
    "# Test functional correctness of quant layer implementation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "2ed08ae2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at bert-base-uncased were not used when initializing BertModel: ['cls.seq_relationship.weight', 'cls.predictions.bias', 'cls.predictions.decoder.weight', 'cls.predictions.transform.dense.weight', 'cls.predictions.transform.dense.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.transform.LayerNorm.bias', 'cls.seq_relationship.bias']\n",
      "- This IS expected if you are initializing BertModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing BertModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from transformers import BertModel, BertTokenizer\n",
    "import math\n",
    "\n",
    "from src.quant_layer import attention, ffn, layer_kernel_gt\n",
    "from src import quant_layer\n",
    "from src.quant_ops import tensor_quant_scale\n",
    "\n",
    "tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')\n",
    "model = BertModel.from_pretrained('bert-base-uncased')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "31dec032",
   "metadata": {},
   "outputs": [],
   "source": [
    "text_512 = 'This project aims to implement a transformer layer on a cluster of FPGAs. In recent years transformers have outperformed traditional convolutional neural networks in many fields, but serial performance is dismal and parallel GPU performance is power-intensive. Specialized architectures have been studied little, especially using FPGA platforms. This research will improve transformer inference performance by offloading computationally intensive sections of the network to reconfigurable accelerators running on a cluster of multiple FPGA devices. This research will result in an acceleration architecture for a single layer of a transformer network along with a performance comparison with CPU and GPU baselines. We propose the investigation of distributed transformer inference across a cluster of multiple field programmable gate arrays (FPGAs). This research will investigate the partitioning of a transformer layer across multiple FPGA devices along with networking between FPGAs in the cluster. Transformers have become a dominant machine learning architecture for many domains such as natural language processing, therefore high speed inference is desirable. However, networks sizes and limited FPGA resources often make inference on a single FPGA slow due to limited parallelism and pipeline depth or impossible due to limited resources. The purpose of this research is to explore methods to overcome these challenges by introducing parallelism through multi-FPGA clusters. Transformers are highly parallel neural network architectures which consist of stacks of encoder and decoder layers. These layers consist of many linear transformations on matrices which are represented by matrix-matrix multiplication. Within an encoder/decoder layer there is an opportunity to parallelize both between concurrent general matrix multiplies (GeMM) and within each GeMM. Attempting to serialize these operations on a CPU leads to high execution time and is a poor utilization of the CPU\\'s general purpose architecture. GPUs can deliver high throughput inference for transformers, though they are power-hungry and do not achieve the low latency required by some applications. Both in the datacenter and at the edge, low-latency and efficient inference is desired. Optimally, there would be an architecture that could scale between these two extremes of computational demand. State-of-the-art transformers can contain upwards of 12 layers and multiply matrices on the order of 1024x1024 elements. In addition, the trend of increasing transformer size does not show signs of slowing. This large use of memory and FLOPs leads to difficulty mapping an entire transformer network to a '\n",
    "text_128 = 'This project aims to implement a transformer layer on a cluster of FPGAs. In recent years transformers have outperformed traditional convolutional neural networks in many fields, but serial performance is dismal and parallel GPU performance is power-intensive. Specialized architectures have been studied little, especially using FPGA platforms. This research will improve transformer inference performance by offloading computationally intensive sections of the network to reconfigurable accelerators running on a cluster of multiple FPGA devices. This research will result in an acceleration architecture for a single layer of a transformer network along with a  '\n",
    "text = text_128\n",
    "encoded_input = tokenizer(text, return_tensors='pt')\n",
    "embedding_output = model.embeddings(\n",
    "    input_ids=encoded_input['input_ids'],\n",
    "    position_ids=None,\n",
    "    token_type_ids=encoded_input['token_type_ids'],\n",
    "    inputs_embeds=None,\n",
    "    past_key_values_length=0,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "fa9789dd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 128, 768])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "embedding_output.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7bb56e01",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "attention_probs_scale 0.007628676470588235\n",
      "attention_probs_int tensor([[[[  1.,   2.,   1.,  ...,   2.,   2.,   3.],\n",
      "          [  4.,   1.,   2.,  ...,   1.,   0.,   1.],\n",
      "          [  2.,   1.,   1.,  ...,   0.,   0.,   1.],\n",
      "          ...,\n",
      "          [  1.,   1.,   1.,  ...,   2.,   1.,   2.],\n",
      "          [  1.,   1.,   1.,  ...,   2.,   2.,   2.],\n",
      "          [  1.,   1.,   1.,  ...,   1.,   2.,   2.]],\n",
      "\n",
      "         [[  7.,   0.,   0.,  ...,   3.,   0.,   0.],\n",
      "          [  0.,   0.,   2.,  ...,   0.,   0.,   0.],\n",
      "          [  0.,   0.,   0.,  ...,   0.,   0.,   1.],\n",
      "          ...,\n",
      "          [  0.,   0.,   2.,  ...,   1.,   3.,   0.],\n",
      "          [  0.,   1.,   1.,  ...,   3.,   3.,   0.],\n",
      "          [  0.,   1.,   0.,  ...,  12.,   4.,   0.]],\n",
      "\n",
      "         [[ 48.,   2.,   1.,  ...,   1.,   1.,   2.],\n",
      "          [ 70.,  12.,   1.,  ...,   0.,   1.,   1.],\n",
      "          [ 13.,  65.,   4.,  ...,   1.,   0.,   1.],\n",
      "          ...,\n",
      "          [  4.,   0.,   0.,  ...,   1.,   2.,   6.],\n",
      "          [ 31.,   4.,   0.,  ...,  35.,   8.,  12.],\n",
      "          [ 34.,   1.,   1.,  ...,   5.,  47.,  17.]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[  1.,   1.,   1.,  ...,   0.,   0.,   0.],\n",
      "          [  6.,   1.,   1.,  ...,   0.,   0.,   1.],\n",
      "          [  2.,   0.,   7.,  ...,   0.,   0.,   1.],\n",
      "          ...,\n",
      "          [  8.,   1.,   2.,  ...,   0.,   0.,   1.],\n",
      "          [ 12.,   1.,   1.,  ...,   1.,   0.,   1.],\n",
      "          [ 19.,   1.,   1.,  ...,   1.,   0.,   1.]],\n",
      "\n",
      "         [[ 44.,   2.,   0.,  ...,   0.,   1.,   1.],\n",
      "          [  0.,   8.,  58.,  ...,   0.,   1.,   2.],\n",
      "          [ 23.,   9.,   5.,  ...,   0.,   0.,   3.],\n",
      "          ...,\n",
      "          [  0.,   0.,   0.,  ...,   4.,  31.,  19.],\n",
      "          [  0.,   0.,   0.,  ...,   6.,   4., 108.],\n",
      "          [  5.,   0.,   0.,  ...,  10.,  17.,  64.]],\n",
      "\n",
      "         [[ 78.,   5.,   0.,  ...,   0.,   1.,   1.],\n",
      "          [ 17.,   7.,   8.,  ...,   0.,   0.,   1.],\n",
      "          [  1.,   4.,   1.,  ...,   1.,   0.,   1.],\n",
      "          ...,\n",
      "          [  4.,   1.,   1.,  ...,   2.,   2.,   3.],\n",
      "          [  5.,   2.,   1.,  ...,   9.,   4.,   7.],\n",
      "          [ 31.,   4.,   1.,  ...,   5.,  20.,  11.]]]],\n",
      "       grad_fn=<RoundBackward0>)\n",
      "attention_out tensor([[[[-29.,   3.,   5.,  ...,   3.,  -5.,  -2.],\n",
      "          [-31.,   5.,   4.,  ...,   3.,   1.,  -5.],\n",
      "          [-29.,   2.,   5.,  ...,  -1.,  -0.,  -7.],\n",
      "          ...,\n",
      "          [-27.,   5.,   3.,  ...,   3.,  -1.,  -1.],\n",
      "          [-29.,   3.,   5.,  ...,   2.,  -4.,  -3.],\n",
      "          [-26.,   5.,   2.,  ...,   5.,  -2.,  -2.]],\n",
      "\n",
      "         [[  1.,  27.,   1.,  ...,   3.,   9.,  -2.],\n",
      "          [  4.,   0., -16.,  ...,  -4.,  -1.,   1.],\n",
      "          [ -1.,   1., -25.,  ...,  11.,  -7.,   4.],\n",
      "          ...,\n",
      "          [  2.,   4., -16.,  ...,  -0.,  -3.,   5.],\n",
      "          [  3.,   8.,  -6.,  ...,  -4.,   1.,   7.],\n",
      "          [ -2.,  19.,   3.,  ...,  -1.,   0.,   2.]],\n",
      "\n",
      "         [[  7.,   0.,  -4.,  ...,  -6.,  -5.,   8.],\n",
      "          [  6.,  -9.,  -3.,  ...,  -4.,  -9.,   2.],\n",
      "          [ -5., -23.,  -8.,  ..., -37., -47.,  -1.],\n",
      "          ...,\n",
      "          [-18., -17.,  -6.,  ..., -12., -13.,  -7.],\n",
      "          [ -5.,   1.,  -4.,  ..., -15.,  -4.,  20.],\n",
      "          [  8., -13.,  -0.,  ..., -13., -10.,  26.]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[ -8.,  -1.,   5.,  ...,   8.,   8.,  -9.],\n",
      "          [ -5.,  -9.,   2.,  ...,   0.,   9.,  -7.],\n",
      "          [-10., -11.,   2.,  ...,  -7.,   7., -10.],\n",
      "          ...,\n",
      "          [ -4.,  -5.,   2.,  ...,  -1.,   9.,  -7.],\n",
      "          [ -5., -10.,   3.,  ...,   2.,   7.,  -7.],\n",
      "          [ -1.,  -9.,   2.,  ...,   0.,   4.,  -6.]],\n",
      "\n",
      "         [[  1.,   1.,   3.,  ...,  -4.,  -2.,   6.],\n",
      "          [ -2.,  34.,   9.,  ...,  28.,   4.,   2.],\n",
      "          [  1.,  15.,   3.,  ...,  -2.,   0.,  10.],\n",
      "          ...,\n",
      "          [  3., -25.,  -1.,  ...,   2.,   0., -15.],\n",
      "          [-29., -58.,  25.,  ...,   5., -36., -24.],\n",
      "          [-13., -31.,  16.,  ...,   1., -23., -16.]],\n",
      "\n",
      "         [[ -5.,  -4.,  -4.,  ...,  -0.,   6.,   3.],\n",
      "          [-10.,  -8.,  -2.,  ...,  -2.,  -4.,  -3.],\n",
      "          [ -6.,  -3.,  -4.,  ..., -13.,   5.,  -1.],\n",
      "          ...,\n",
      "          [-11.,  -9.,  -4.,  ...,  -6.,   3.,   3.],\n",
      "          [-16.,  -9.,  -3.,  ...,  -0.,   0.,  -1.],\n",
      "          [-14.,  -5.,  -9.,  ...,   1.,   3.,   2.]]]],\n",
      "       grad_fn=<RoundBackward0>)\n",
      "dense_out tensor([[[ 0.0733, -0.2733, -0.6244,  ...,  0.0997,  0.2117,  0.2552],\n",
      "         [-0.7298,  0.7013,  0.1123,  ...,  0.4356,  0.8463,  0.2713],\n",
      "         [ 0.2087,  0.3795,  0.0258,  ...,  0.0904,  0.0506, -0.5959],\n",
      "         ...,\n",
      "         [-0.8258,  0.8407,  0.9534,  ...,  0.1572, -1.1183,  0.9325],\n",
      "         [ 0.1323,  0.6437,  0.5281,  ...,  0.0379,  0.0143,  0.5286],\n",
      "         [-0.3892,  0.4597,  0.2496,  ..., -0.7265, -0.1934, -0.1269]]],\n",
      "       grad_fn=<AddBackward0>)\n",
      "tensor([[[ 0.1090, -0.0411, -0.1556,  ...,  0.0649,  0.0122,  0.0245],\n",
      "         [-0.4882,  0.4328,  0.1707,  ...,  0.0806,  0.5541,  0.1315],\n",
      "         [ 0.2778,  0.4525,  0.1470,  ..., -0.0282,  0.3554, -0.8637],\n",
      "         ...,\n",
      "         [-0.4865,  0.6699,  0.7012,  ...,  0.0566, -1.1085,  0.7002],\n",
      "         [ 0.3689,  0.5432,  0.7258,  ...,  0.3103, -0.0265,  0.6428],\n",
      "         [-0.1264,  0.1299,  0.1691,  ..., -0.1081,  0.0902, -0.3762]]],\n",
      "       grad_fn=<MulBackward0>)\n"
     ]
    }
   ],
   "source": [
    "layer = model.encoder.layer[0]\n",
    "attention_out = attention(layer, embedding_output)\n",
    "output_gt = ffn(layer, attention_out)\n",
    "print(output_gt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3c137792",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[  1,  -1,  -2,  ...,   1,   0,   0],\n",
       "         [ -6,   6,   2,  ...,   1,   7,   2],\n",
       "         [  4,   6,   2,  ...,   0,   5, -11],\n",
       "         ...,\n",
       "         [ -6,   9,   9,  ...,   1, -14,   9],\n",
       "         [  5,   7,   9,  ...,   4,   0,   8],\n",
       "         [ -2,   2,   2,  ...,  -1,   1,  -5]]], dtype=torch.int8)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output_gt_int, output_scale = tensor_quant_scale(output_gt, bits=8)\n",
    "output_gt_int = output_gt_int.type(torch.int8)\n",
    "output_gt_int"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "5837757b",
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "pipeline() missing 2 required positional arguments: 'stage3_args' and 'stage4_args'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "Input \u001b[0;32mIn [6]\u001b[0m, in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0m output_test \u001b[38;5;241m=\u001b[39m \u001b[43mquant_layer\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mpipeline\u001b[49m\u001b[43m(\u001b[49m\u001b[43mlayer\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43membedding_output\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m      2\u001b[0m \u001b[38;5;28mprint\u001b[39m(output_test)\n",
      "\u001b[0;31mTypeError\u001b[0m: pipeline() missing 2 required positional arguments: 'stage3_args' and 'stage4_args'"
     ]
    }
   ],
   "source": [
    "output_test = quant_layer.pipeline(layer, embedding_output)\n",
    "print(output_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ff18794c",
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'output_test' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Input \u001b[0;32mIn [7]\u001b[0m, in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m torch\u001b[38;5;241m.\u001b[39mallclose(output_gt, \u001b[43moutput_test\u001b[49m)\n",
      "\u001b[0;31mNameError\u001b[0m: name 'output_test' is not defined"
     ]
    }
   ],
   "source": [
    "assert torch.allclose(output_gt, output_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e248a2d",
   "metadata": {},
   "source": [
    "# Test stage by stage"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ca7c28b",
   "metadata": {},
   "source": [
    "## Stage 1\n",
    "Quantize the ground truth output so we can compare it to the quantized output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f9646dc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "(stage1_query_gt, stage1_key_gt, stage1_value_gt) = quant_layer.stage1_dynamic(layer, embedding_output)\n",
    "\n",
    "stage1_query_gt_int, _ = tensor_quant_scale(stage1_query_gt)\n",
    "stage1_query_gt_int = stage1_query_gt_int.type(torch.int8)\n",
    "stage1_key_gt_int, _ = tensor_quant_scale(stage1_key_gt)\n",
    "stage1_key_gt_int = stage1_key_gt_int.type(torch.int8)\n",
    "stage1_value_gt_int, _ = tensor_quant_scale(stage1_value_gt)\n",
    "stage1_value_gt_int = stage1_value_gt_int.type(torch.int8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "ba7693e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "stage1_args, stage2_args, stage3_args, stage4_args = layer_kernel_gt(layer, embedding_output)\n",
    "stage1_query_uut, stage1_key_uut, stage1_value_uut = quant_layer.stage1(**stage1_args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "fc6a91c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert torch.allclose(stage1_query_gt_int, stage1_query_uut)\n",
    "assert torch.allclose(stage1_key_gt_int, stage1_key_uut)\n",
    "assert torch.allclose(stage1_value_gt_int, stage1_value_uut)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9196e982",
   "metadata": {},
   "source": [
    "## Stage 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "38fa8b3c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "attention_probs_scale 0.007628676470588235\n",
      "attention_probs_int tensor([[[[  1.,   2.,   1.,  ...,   2.,   2.,   3.],\n",
      "          [  4.,   1.,   2.,  ...,   1.,   0.,   1.],\n",
      "          [  2.,   1.,   1.,  ...,   0.,   0.,   1.],\n",
      "          ...,\n",
      "          [  1.,   1.,   1.,  ...,   2.,   1.,   2.],\n",
      "          [  1.,   1.,   1.,  ...,   2.,   2.,   2.],\n",
      "          [  1.,   1.,   1.,  ...,   1.,   2.,   2.]],\n",
      "\n",
      "         [[  7.,   0.,   0.,  ...,   3.,   0.,   0.],\n",
      "          [  0.,   0.,   2.,  ...,   0.,   0.,   0.],\n",
      "          [  0.,   0.,   0.,  ...,   0.,   0.,   1.],\n",
      "          ...,\n",
      "          [  0.,   0.,   2.,  ...,   1.,   3.,   0.],\n",
      "          [  0.,   1.,   1.,  ...,   3.,   3.,   0.],\n",
      "          [  0.,   1.,   0.,  ...,  12.,   4.,   0.]],\n",
      "\n",
      "         [[ 48.,   2.,   1.,  ...,   1.,   1.,   2.],\n",
      "          [ 70.,  12.,   1.,  ...,   0.,   1.,   1.],\n",
      "          [ 13.,  65.,   4.,  ...,   1.,   0.,   1.],\n",
      "          ...,\n",
      "          [  4.,   0.,   0.,  ...,   1.,   2.,   6.],\n",
      "          [ 31.,   4.,   0.,  ...,  35.,   8.,  12.],\n",
      "          [ 34.,   1.,   1.,  ...,   5.,  47.,  17.]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[  1.,   1.,   1.,  ...,   0.,   0.,   0.],\n",
      "          [  6.,   1.,   1.,  ...,   0.,   0.,   1.],\n",
      "          [  2.,   0.,   7.,  ...,   0.,   0.,   1.],\n",
      "          ...,\n",
      "          [  8.,   1.,   2.,  ...,   0.,   0.,   1.],\n",
      "          [ 12.,   1.,   1.,  ...,   1.,   0.,   1.],\n",
      "          [ 19.,   1.,   1.,  ...,   1.,   0.,   1.]],\n",
      "\n",
      "         [[ 44.,   2.,   0.,  ...,   0.,   1.,   1.],\n",
      "          [  0.,   8.,  58.,  ...,   0.,   1.,   2.],\n",
      "          [ 23.,   9.,   5.,  ...,   0.,   0.,   3.],\n",
      "          ...,\n",
      "          [  0.,   0.,   0.,  ...,   4.,  31.,  19.],\n",
      "          [  0.,   0.,   0.,  ...,   6.,   4., 108.],\n",
      "          [  5.,   0.,   0.,  ...,  10.,  17.,  64.]],\n",
      "\n",
      "         [[ 78.,   5.,   0.,  ...,   0.,   1.,   1.],\n",
      "          [ 17.,   7.,   8.,  ...,   0.,   0.,   1.],\n",
      "          [  1.,   4.,   1.,  ...,   1.,   0.,   1.],\n",
      "          ...,\n",
      "          [  4.,   1.,   1.,  ...,   2.,   2.,   3.],\n",
      "          [  5.,   2.,   1.,  ...,   9.,   4.,   7.],\n",
      "          [ 31.,   4.,   1.,  ...,   5.,  20.,  11.]]]],\n",
      "       grad_fn=<RoundBackward0>)\n",
      "attention_out tensor([[[[-29.,   3.,   5.,  ...,   3.,  -5.,  -2.],\n",
      "          [-31.,   5.,   4.,  ...,   3.,   1.,  -5.],\n",
      "          [-29.,   2.,   5.,  ...,  -1.,  -0.,  -7.],\n",
      "          ...,\n",
      "          [-27.,   5.,   3.,  ...,   3.,  -1.,  -1.],\n",
      "          [-29.,   3.,   5.,  ...,   2.,  -4.,  -3.],\n",
      "          [-26.,   5.,   2.,  ...,   5.,  -2.,  -2.]],\n",
      "\n",
      "         [[  1.,  27.,   1.,  ...,   3.,   9.,  -2.],\n",
      "          [  4.,   0., -16.,  ...,  -4.,  -1.,   1.],\n",
      "          [ -1.,   1., -25.,  ...,  11.,  -7.,   4.],\n",
      "          ...,\n",
      "          [  2.,   4., -16.,  ...,  -0.,  -3.,   5.],\n",
      "          [  3.,   8.,  -6.,  ...,  -4.,   1.,   7.],\n",
      "          [ -2.,  19.,   3.,  ...,  -1.,   0.,   2.]],\n",
      "\n",
      "         [[  7.,   0.,  -4.,  ...,  -6.,  -5.,   8.],\n",
      "          [  6.,  -9.,  -3.,  ...,  -4.,  -9.,   2.],\n",
      "          [ -5., -23.,  -8.,  ..., -37., -47.,  -1.],\n",
      "          ...,\n",
      "          [-18., -17.,  -6.,  ..., -12., -13.,  -7.],\n",
      "          [ -5.,   1.,  -4.,  ..., -15.,  -4.,  20.],\n",
      "          [  8., -13.,  -0.,  ..., -13., -10.,  26.]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[ -8.,  -1.,   5.,  ...,   8.,   8.,  -9.],\n",
      "          [ -5.,  -9.,   2.,  ...,   0.,   9.,  -7.],\n",
      "          [-10., -11.,   2.,  ...,  -7.,   7., -10.],\n",
      "          ...,\n",
      "          [ -4.,  -5.,   2.,  ...,  -1.,   9.,  -7.],\n",
      "          [ -5., -10.,   3.,  ...,   2.,   7.,  -7.],\n",
      "          [ -1.,  -9.,   2.,  ...,   0.,   4.,  -6.]],\n",
      "\n",
      "         [[  1.,   1.,   3.,  ...,  -4.,  -2.,   6.],\n",
      "          [ -2.,  34.,   9.,  ...,  28.,   4.,   2.],\n",
      "          [  1.,  15.,   3.,  ...,  -2.,   0.,  10.],\n",
      "          ...,\n",
      "          [  3., -25.,  -1.,  ...,   2.,   0., -15.],\n",
      "          [-29., -58.,  25.,  ...,   5., -36., -24.],\n",
      "          [-13., -31.,  16.,  ...,   1., -23., -16.]],\n",
      "\n",
      "         [[ -5.,  -4.,  -4.,  ...,  -0.,   6.,   3.],\n",
      "          [-10.,  -8.,  -2.,  ...,  -2.,  -4.,  -3.],\n",
      "          [ -6.,  -3.,  -4.,  ..., -13.,   5.,  -1.],\n",
      "          ...,\n",
      "          [-11.,  -9.,  -4.,  ...,  -6.,   3.,   3.],\n",
      "          [-16.,  -9.,  -3.,  ...,  -0.,   0.,  -1.],\n",
      "          [-14.,  -5.,  -9.,  ...,   1.,   3.,   2.]]]],\n",
      "       grad_fn=<RoundBackward0>)\n",
      "dense_out tensor([[[ 0.0733, -0.2733, -0.6244,  ...,  0.0997,  0.2117,  0.2552],\n",
      "         [-0.7298,  0.7013,  0.1123,  ...,  0.4356,  0.8463,  0.2713],\n",
      "         [ 0.2087,  0.3795,  0.0258,  ...,  0.0904,  0.0506, -0.5959],\n",
      "         ...,\n",
      "         [-0.8258,  0.8407,  0.9534,  ...,  0.1572, -1.1183,  0.9325],\n",
      "         [ 0.1323,  0.6437,  0.5281,  ...,  0.0379,  0.0143,  0.5286],\n",
      "         [-0.3892,  0.4597,  0.2496,  ..., -0.7265, -0.1934, -0.1269]]],\n",
      "       grad_fn=<AddBackward0>)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[[ 1, -1, -2,  ...,  1,  0,  1],\n",
       "         [-2,  2, -1,  ...,  2,  2,  0],\n",
       "         [ 1,  1, -1,  ...,  1,  0, -2],\n",
       "         ...,\n",
       "         [-2,  2,  3,  ...,  0, -4,  3],\n",
       "         [ 1,  2,  1,  ...,  0,  0,  2],\n",
       "         [-1,  2,  1,  ..., -2, -1, -1]]], dtype=torch.int8)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "stage2_gt = quant_layer.stage2_dynamic(layer, embedding_output, stage1_query_gt, stage1_key_gt, stage1_value_gt)\n",
    "stage2_gt_int, _ = tensor_quant_scale(stage2_gt, bits=8)\n",
    "stage2_gt_int = stage2_gt_int.type(torch.int8)\n",
    "stage2_gt_int"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "2c33b05d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "attention_probs_int tensor([[[[  1,   1,   0,  ...,   1,   1,   2],\n",
      "          [  3,   0,   1,  ...,   0,   0,   0],\n",
      "          [  1,   0,   0,  ...,   0,   0,   0],\n",
      "          ...,\n",
      "          [  1,   0,   0,  ...,   2,   1,   2],\n",
      "          [  0,   0,   0,  ...,   1,   2,   1],\n",
      "          [  1,   1,   0,  ...,   1,   1,   2]],\n",
      "\n",
      "         [[  6,   0,   0,  ...,   3,   0,   0],\n",
      "          [  0,   0,   2,  ...,   0,   0,   0],\n",
      "          [  0,   0,   0,  ...,   0,   0,   0],\n",
      "          ...,\n",
      "          [  0,   0,   1,  ...,   0,   2,   0],\n",
      "          [  0,   0,   0,  ...,   2,   3,   0],\n",
      "          [  0,   0,   0,  ...,  12,   4,   0]],\n",
      "\n",
      "         [[ 47,   1,   0,  ...,   0,   0,   1],\n",
      "          [ 68,  11,   1,  ...,   0,   0,   1],\n",
      "          [ 13,  63,   3,  ...,   1,   0,   1],\n",
      "          ...,\n",
      "          [  4,   0,   0,  ...,   0,   2,   5],\n",
      "          [ 30,   3,   0,  ...,  34,   8,  12],\n",
      "          [ 33,   1,   0,  ...,   5,  45,  16]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[  1,   1,   0,  ...,   0,   0,   0],\n",
      "          [  6,   0,   1,  ...,   0,   0,   0],\n",
      "          [  2,   0,   6,  ...,   0,   0,   0],\n",
      "          ...,\n",
      "          [  8,   0,   1,  ...,   0,   0,   0],\n",
      "          [ 11,   0,   0,  ...,   0,   0,   1],\n",
      "          [ 18,   0,   1,  ...,   0,   0,   0]],\n",
      "\n",
      "         [[ 43,   1,   0,  ...,   0,   0,   0],\n",
      "          [  0,   8,  57,  ...,   0,   0,   1],\n",
      "          [ 22,   9,   5,  ...,   0,   0,   2],\n",
      "          ...,\n",
      "          [  0,   0,   0,  ...,   3,  30,  19],\n",
      "          [  0,   0,   0,  ...,   5,   4, 105],\n",
      "          [  4,   0,   0,  ...,   9,  17,  62]],\n",
      "\n",
      "         [[ 76,   4,   0,  ...,   0,   1,   1],\n",
      "          [ 17,   6,   7,  ...,   0,   0,   1],\n",
      "          [  0,   3,   0,  ...,   0,   0,   0],\n",
      "          ...,\n",
      "          [  3,   0,   1,  ...,   2,   2,   3],\n",
      "          [  4,   1,   1,  ...,   8,   4,   7],\n",
      "          [ 30,   4,   0,  ...,   4,  19,  11]]]], dtype=torch.int8)\n",
      "attention_out tensor([[[[ -4,   0,   0,  ...,   2,  -2,   1],\n",
      "          [ -6,   2,   0,  ...,   1,   1,  -1],\n",
      "          [ -7,   0,   1,  ...,   0,   1,  -3],\n",
      "          ...,\n",
      "          [ -4,   2,   0,  ...,   1,   0,   0],\n",
      "          [ -5,   1,   1,  ...,   1,  -1,   0],\n",
      "          [ -6,   1,   0,  ...,   1,  -1,  -1]],\n",
      "\n",
      "         [[  1,  13,   0,  ...,   2,   4,  -1],\n",
      "          [  2,  -1,  -5,  ...,  -3,   0,   1],\n",
      "          [  0,   0, -10,  ...,   5,  -4,   2],\n",
      "          ...,\n",
      "          [  1,   1,  -6,  ...,  -1,  -1,   2],\n",
      "          [  1,   3,  -2,  ...,  -2,   1,   3],\n",
      "          [ -1,   8,   2,  ...,   0,   0,   0]],\n",
      "\n",
      "         [[  1,   0,  -1,  ...,   0,   0,   1],\n",
      "          [  3,  -4,  -2,  ...,  -2,  -4,   0],\n",
      "          [ -3, -12,  -4,  ..., -19, -23,  -1],\n",
      "          ...,\n",
      "          [ -9,  -9,  -3,  ...,  -6,  -6,  -4],\n",
      "          [ -3,   0,  -2,  ...,  -7,  -2,   8],\n",
      "          [  3,  -6,   0,  ...,  -6,  -5,  12]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[ -2,   0,   2,  ...,   3,   3,  -5],\n",
      "          [ -1,  -2,   0,  ...,  -1,   2,  -2],\n",
      "          [ -5,  -4,   1,  ...,  -4,   1,  -4],\n",
      "          ...,\n",
      "          [ -1,  -1,   1,  ...,   0,   1,  -2],\n",
      "          [ -1,  -4,   1,  ...,   1,   1,  -4],\n",
      "          [  1,  -2,   1,  ...,   1,   0,  -3]],\n",
      "\n",
      "         [[  0,  -1,   1,  ...,   0,  -2,   2],\n",
      "          [ -1,  16,   4,  ...,  14,   2,   1],\n",
      "          [  0,   7,   1,  ...,  -1,   1,   5],\n",
      "          ...,\n",
      "          [  1, -13,  -1,  ...,   1,   0,  -8],\n",
      "          [-14, -29,  12,  ...,   2, -18, -12],\n",
      "          [ -7, -16,   8,  ...,   0, -11,  -8]],\n",
      "\n",
      "         [[ -2,  -2,  -2,  ...,   0,   2,   2],\n",
      "          [ -4,  -4,  -1,  ...,   0,  -2,  -1],\n",
      "          [  0,   0,  -3,  ...,  -5,   2,   0],\n",
      "          ...,\n",
      "          [ -3,  -3,  -2,  ...,  -3,   2,   3],\n",
      "          [ -5,  -3,  -2,  ...,   0,   1,   1],\n",
      "          [ -5,  -2,  -4,  ...,   0,   2,   1]]]], dtype=torch.int8)\n",
      "dense_out tensor([[[ 0.0535, -0.2364, -0.5940,  ...,  0.0224,  0.1786,  0.2165],\n",
      "         [-0.6960,  0.7347,  0.1651,  ...,  0.3337,  0.8397,  0.2611],\n",
      "         [ 0.2382,  0.4259,  0.0856,  ...,  0.0212,  0.0404, -0.6309],\n",
      "         ...,\n",
      "         [-0.7827,  0.8944,  1.0181,  ...,  0.0399, -1.0734,  0.8836],\n",
      "         [ 0.1514,  0.7107,  0.5943,  ..., -0.0344,  0.0389,  0.5022],\n",
      "         [-0.3723,  0.4842,  0.3034,  ..., -0.7702, -0.1804, -0.1615]]],\n",
      "       grad_fn=<AddBackward0>)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[[  2,  -2,  -7,  ...,   1,   1,   2],\n",
       "         [ -6,   7,   0,  ...,   4,   8,   2],\n",
       "         [  4,   4,  -1,  ...,   1,   0,  -7],\n",
       "         ...,\n",
       "         [ -7,   8,   8,  ...,   1, -11,   8],\n",
       "         [  3,   7,   4,  ...,   0,  -1,   4],\n",
       "         [ -3,   5,   2,  ...,  -7,  -3,  -3]]], dtype=torch.int8)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "stage2_uut = quant_layer.stage2(stage1_query_uut, stage1_key_uut, stage1_value_uut, **stage2_args)\n",
    "stage2_uut"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "fc314b78",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.00784313725490196"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(1/stage2_args['M_attention_probs'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "019054f0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.0070562844355619315, 0.008036052309477855, 0.003417007858954613)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "stage1_args['M_query'], stage1_args['M_key'], stage1_args['M_value'] "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4d66d75",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9fdc999e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "b9433b41",
   "metadata": {},
   "source": [
    "Testing C++ impl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "1f6777fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "def genmat(A, mod):\n",
    "    for i in range(A.size()[0]):\n",
    "        for j in range(A.size()[1]):\n",
    "            A[i][j] = (i*dmodel+j) % mod"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "f5e22b67",
   "metadata": {},
   "outputs": [],
   "source": [
    "seqlen = 6\n",
    "dmodel = 8\n",
    "hidden_states = torch.zeros(seqlen,dmodel)\n",
    "query_weight_t = torch.zeros(dmodel,dmodel)\n",
    "key_weight_t = torch.zeros(dmodel,dmodel)\n",
    "value_weight_t = torch.zeros(dmodel,dmodel)\n",
    "query_bias = torch.zeros(1,dmodel)\n",
    "key_bias = torch.zeros(1,dmodel)\n",
    "value_bias = torch.zeros(1,dmodel)\n",
    "genmat(hidden_states, 7)\n",
    "genmat(query_weight_t, 9)\n",
    "genmat(key_weight_t, 11)\n",
    "genmat(value_weight_t, 13)\n",
    "genmat(query_bias, 63)\n",
    "genmat(key_bias, 65)\n",
    "genmat(value_bias, 67)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "eca172c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "query = torch.matmul(hidden_states, query_weight_t) + query_bias\n",
    "query *= 0.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "2f993cde",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[49.0000, 55.5000, 57.5000, 55.0000, 48.0000, 36.5000, 20.5000, 31.5000],\n",
       "        [56.0000, 58.5000, 56.5000, 50.0000, 39.0000, 23.5000, 35.0000, 42.0000],\n",
       "        [59.5000, 58.0000, 52.0000, 41.5000, 26.5000, 38.5000, 46.0000, 49.0000],\n",
       "        [59.5000, 54.0000, 44.0000, 29.5000, 42.0000, 50.0000, 53.5000, 52.5000],\n",
       "        [56.0000, 46.5000, 32.5000, 45.5000, 54.0000, 58.0000, 57.5000, 52.5000],\n",
       "        [49.0000, 35.5000, 49.0000, 58.0000, 62.5000, 62.5000, 58.0000, 49.0000]])"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "92555b04",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "5d0003a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_in = torch.Tensor(np.arange(24))\n",
    "key_in = torch.Tensor(np.arange(24))\n",
    "query_in = query_in.view((2, 3, 4))\n",
    "key_in = key_in.view((2, 3, 4))\n",
    "query_in = query_in.permute(1,0,2)\n",
    "key_in = key_in.permute(1,2,0)\n",
    "att_scores = torch.matmul(query_in, key_in)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "9105947e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([  14,   86,   86,  734,  126,  390,  390, 1230,  366,  822,  822, 1854],\n",
       "       dtype=torch.int32)"
      ]
     },
     "execution_count": 52,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "att_scores.type(torch.int).flatten()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "0439fbae",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16,\n",
       "       17, 18, 19, 20, 21, 22, 23])"
      ]
     },
     "execution_count": 53,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.arange(24)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "6bc48f3d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([ 0.,  1.,  2.,  3., 12., 13., 14., 15.,  4.,  5.,  6.,  7., 16., 17.,\n",
       "        18., 19.,  8.,  9., 10., 11., 20., 21., 22., 23.])"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query_in.flatten()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "555421f7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([ 0., 12.,  1., 13.,  2., 14.,  3., 15.,  4., 16.,  5., 17.,  6., 18.,\n",
       "         7., 19.,  8., 20.,  9., 21., 10., 22., 11., 23.])"
      ]
     },
     "execution_count": 55,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "key_in.contiguous().flatten()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "d67c550b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([3, 4, 2])"
      ]
     },
     "execution_count": 56,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "key_in.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "f04e55a8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 0., 12.],\n",
       "         [ 1., 13.],\n",
       "         [ 2., 14.],\n",
       "         [ 3., 15.]],\n",
       "\n",
       "        [[ 4., 16.],\n",
       "         [ 5., 17.],\n",
       "         [ 6., 18.],\n",
       "         [ 7., 19.]],\n",
       "\n",
       "        [[ 8., 20.],\n",
       "         [ 9., 21.],\n",
       "         [10., 22.],\n",
       "         [11., 23.]]])"
      ]
     },
     "execution_count": 57,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "key_in"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "d740fca5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([3, 2, 4]), torch.Size([3, 4, 2]))"
      ]
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query_in.shape, key_in.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "id": "b19d5e0a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([ 40.,  43.,  46.,  49., 112., 124., 136., 148., 184., 205., 226., 247.,\n",
       "        376., 406., 436., 466., 484., 523., 562., 601., 592., 640., 688., 736.])"
      ]
     },
     "execution_count": 65,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nhead = 2\n",
    "seqlen = 3\n",
    "dhead = 4\n",
    "probs_in = torch.Tensor(np.arange(nhead*seqlen*seqlen))\n",
    "probs_in = probs_in.view((nhead, seqlen, seqlen))\n",
    "value_in = torch.Tensor(np.arange(nhead*seqlen*dhead))\n",
    "value_in = value_in.view((seqlen, nhead, dhead))\n",
    "value_in = value_in.permute(1,0,2)\n",
    "att_out = torch.matmul(probs_in, value_in)\n",
    "att_out.flatten()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "id": "cd59c8a1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 40.,  43.,  46.,  49., 376., 406., 436., 466.],\n",
       "        [112., 124., 136., 148., 484., 523., 562., 601.],\n",
       "        [184., 205., 226., 247., 592., 640., 688., 736.]])"
      ]
     },
     "execution_count": 66,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "att_out = att_out.permute(1,0,2).contiguous()\n",
    "att_out = att_out.view(seqlen, nhead*dhead)\n",
    "att_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "id": "c6c990a1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([ 40,  43,  46,  49, 376, 406, 436, 466, 112, 124, 136, 148, 484, 523,\n",
       "        562, 601, 184, 205, 226, 247, 592, 640, 688, 736], dtype=torch.int32)"
      ]
     },
     "execution_count": 68,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "att_out.flatten().type(torch.int)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "trans-fat",
   "language": "python",
   "name": "trans-fat"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
